#include <algorithm>
#include <limits>

#include "txdnn_internal.h"
#include "tensor/tensor_descriptor.h"


namespace txdnn {

bool IsDataTypeSupported(txdnnDataType_t t) {
  switch (t) {
    case TXDNN_DATA_FLOAT:
    case TXDNN_DATA_DOUBLE:
    case TXDNN_DATA_HALF:
    case TXDNN_DATA_INT8:
    case TXDNN_DATA_INT32:
    case TXDNN_DATA_UINT8:
    case TXDNN_DATA_BFLOAT16:
    case TXDNN_DATA_INT64:
    case TXDNN_DATA_FP8_E4M3:
    case TXDNN_DATA_FP8_E5M2:
    case TXDNN_DATA_FAST_FLOAT_FOR_FP8: return true;
  }
  return false;
}


bool IsLayoutSupported(txdnnBackendLayoutType_t layout, unsigned num_dims) {
  switch (layout) {
    case TXDNN_LAYOUT_TYPE_PREFERRED_NCHW: return num_dims == 4;
    case TXDNN_LAYOUT_TYPE_PREFERRED_NHWC: return num_dims == 4;
    case TXDNN_LAYOUT_TYPE_PREFERRED_PAD4CK: return num_dims == 4;
    case TXDNN_LAYOUT_TYPE_PREFERRED_PAD8CK: return num_dims == 5;
    case TXDNN_LAYOUT_TYPE_COUNT: return num_dims == 4;
  }
  // TODO
  return false;
}


template <class T>
bool CheckLengths(const std::vector<T>& dims, T max_val = 0) {
  if (dims.empty()) {
    return false;
  }
  if (!std::all_of(dims.cbegin(), dims.cend(), [](T x) { return x > 0; })) {
    return false;
  }
  if (max_val) {
    if (!std::all_of(dims.cbegin(), dims.cend(), [max_val](T x) { return x <= max_val; })) {
      return false;
    }
  }
  return true;
}


void TensorDescriptor::checkArgsAndInit(bool use_strides) {
  if (!IsDataTypeSupported(type_)) {
    // TXDNN_THROW(TXDNN_STATUS_BAD_PARAM, "Unsupported DataType");
    TXDNN_THROW("Unsupported DataType");
  }
  if (dims_.empty()) {
    // TXDNN_THROW(TXDNN_STATUS_BAD_PARAM, "Dimension must be > 1");
    TXDNN_THROW("Dimension must be > 1");
  }
  if (layout_ && !IsLayoutSupported(layout_.value(), dims_.size())) {
    // TXDNN_THROW(TXDNN_STATUS_BAD_PARAM, "Unsupported Layout");
    TXDNN_THROW("Unsupported Layout");
  }
  if (CheckLengths(dims_, static_cast<std::size_t>(std::numeric_limits<int64_t>::max()))) {
    // TXDNN_THROW(TXDNN_STATUS_BAD_PARAM, "Dimension must be > 0 and <= INT64_MAX");
    TXDNN_THROW("Dimension must be > 0 and <= INT64_MAX");
  }
  // TODO
}


std::optional<txdnnBackendLayoutType_t> GetDefaultLayout(unsigned dims) {
  switch(dims) {
    case 4: return txdnnBackendLayoutType_t::TXDNN_LAYOUT_TYPE_PREFERRED_NCHW;
    case 5: return txdnnBackendLayoutType_t::TXDNN_LAYOUT_TYPE_PREFERRED_NHWC;
    default: return std::nullopt;
  }
}


std::vector<std::size_t> ConvertDimsOrThrow(const std::vector<int>& dims, const std::string& msg) {
  if (!CheckLengths(dims)) {
    // TXDNN_THROW(TXDNN_STATUS_BAD_PARAM, msg);
    TXDNN_THROW(msg);
  }
  std::vector<std::size_t> d(dims.cbegin(), dims.cend());
  return d;
}


/* TensorDescriptor Constructor */
TensorDescriptor::TensorDescriptor(txdnnDataType_t t, const std::vector<int>& dims)
  : TensorDescriptor(t, GetDefaultLayout(dims.size()), ConvertDimsOrThrow(dims, "Dims must be >0"), {}, false) {}

TensorDescriptor::TensorDescriptor(txdnnDataType_t t, const std::vector<std::size_t>& dims)
  : TensorDescriptor(t, GetDefaultLayout(dims.size()), dims, {}, false) {}

TensorDescriptor::TensorDescriptor(txdnnDataType_t t, std::vector<std::size_t>&& dims)
  : TensorDescriptor(t, GetDefaultLayout(dims.size()), std::move(dims), {}, false) {}

TensorDescriptor::TensorDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const std::vector<int>& dims)
  : TensorDescriptor(t, layout, ConvertDimsOrThrow(dims, "Dims must be >0")) {}

TensorDescriptor::TensorDescriptor(txdnnDataType_t t, const std::vector<int>& dims, const std::vector<int>& strides)
  : TensorDescriptor(t, ConvertDimsOrThrow(dims, "Dims must be >0"),
                      ConvertDimsOrThrow(strides, "strides must be >0")) {}


TensorDescriptor TensorDescriptor::makeDescriptor(txdnnDataType_t t, const int* dims, int size) {
  if (dims == nullptr || size <= 0) {
    // TXDNN_THROW(TXDNN_STATUS_BAD_PARAM_NULL_POINTER, "null dims pointer or size must be > 0.");
    TXDNN_THROW("null dims pointer or size must be > 0.");
  }
  return {t, std::vector<int>(dims, dims + size)};
}

TensorDescriptor TensorDescriptor::makeDescriptor(txdnnDataType_t t, const std::size_t* dims, int size) {
  if (dims == nullptr || size <= 0) {
    TXDNN_THROW("null dims pointer or size must be > 0.");
  }
  return {t, std::vector<std::size_t>(dims, dims + size)};
}

TensorDescriptor TensorDescriptor::makeDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout, const int* dims,
                                                  int size) {
  if (dims == nullptr || size <= 0) {
    TXDNN_THROW("null dims pointer or size must be > 0.");
  }
  return {t, layout, std::vector<int>(dims, dims + size)};
}

TensorDescriptor TensorDescriptor::makeDescriptor(txdnnDataType_t t, txdnnBackendLayoutType_t layout,
                                                  const std::size_t* dims, int size) {
  if (dims == nullptr || size <= 0) {
    TXDNN_THROW("null dims pointer or size must be > 0.");
  }
  return {t, layout, std::vector<std::size_t>(dims, dims + size)};
}

TensorDescriptor TensorDescriptor::makeDescriptor(txdnnDataType_t t, const int* dims, const int* strides, int size) {
  if (dims == nullptr || strides == nullptr || size <= 0) {
    TXDNN_THROW("null dims pointer or size must be > 0.");
  }
  return {t, std::vector<int>(dims, dims + size), std::vector<int>(strides, strides + size)};
}

TensorDescriptor TensorDescriptor::makeDescriptor(txdnnDataType_t t, const std::size_t* dims,
                                                  const std::size_t* strides, int size) {
  if (dims == nullptr || strides == nullptr || size <= 0) {
    TXDNN_THROW("null dims pointer or size must be > 0.");
  }
  return {t, std::vector<std::size_t>(dims, dims + size), std::vector<std::size_t>(strides, strides + size)};
}

} // namespace txdnn